import os, sys
import SimpleITK as sitk
from SimpleITK import WriteImage
import numpy as np

from PyQt5.QtWidgets import (QWidget, QMessageBox, QTextEdit, QLabel, QAbstractItemView,
	QPushButton, QApplication, QGridLayout, QLineEdit, QFileDialog, QListView)
from PyQt5 import QtWidgets, QtGui
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import QStringListModel

# ##***格式转换工具***## #
class FormatTookit(QWidget):
	def __init__(self):
		super().__init__()

		btn_open = QPushButton(text='打开文件夹')#QLabel('Title')
		self.lbl_open = QLabel('')#QLineEdit()
		btn_open.clicked.connect(self.showDialog)

		btn_save = QPushButton(text='保存nii文件')#QLabel('Author')
		self.edt_save = QLineEdit(os.path.abspath(os.path.curdir))
		btn_save.clicked.connect(self.save_niis)

		img_name = QLabel('选中图片列表')
		self.img_list = QListView()#QTextEdit('请选择图片')
		self.img_list.setEditTriggers(QAbstractItemView.NoEditTriggers)#QAbstractItemView.NoEditTriggers
 
		grid = QGridLayout()
		grid.setSpacing(10) # 创建了一个网格布局并且设置了组件之间的间距。
 
		grid.addWidget(btn_open, 1, 0)
		grid.addWidget(self.lbl_open, 1, 1)
 
		grid.addWidget(btn_save, 2, 0)
		grid.addWidget(self.edt_save, 2, 1)
 
		grid.addWidget(img_name, 3, 0)
		grid.addWidget(self.img_list, 3, 1, 5, 1) # 如果我们向网格布局中增加一个组件，我们可以提供组件的跨行
		# 和跨列参数。在这个例子中，我们让reviewEdit组件跨了5行。


		# print(self.edt_save.text())
		self.setLayout(grid)
 
		self.setGeometry(300, 300, 600, 300)#x,y,w,h
		self.setWindowTitle('二维图像打包工具(UESTC_TYB@https://gitee.com/tyb311/img-pack-tool)')

		# 设置窗口图标
		ico_path = os.path.join('bag64.ico')
		icon = QtGui.QIcon()
		icon.addPixmap(QtGui.QPixmap(ico_path), QtGui.QIcon.Normal, QtGui.QIcon.Off)
		self.setWindowIcon(icon)
		self.show()

	def showDialog(self):
		# 弹出文件选择框。第一个字符串参数是getOpenFileName()方法的标题。第二个字符串参数指定了对话框的工作目录。
		# 默认的，文件过滤器设置成All files (*)。
		folder_dir = QFileDialog.getExistingDirectory(self, '打开文件夹', os.path.abspath(os.path.curdir))
		# folder_dir = QFileDialog.getExistingDirectory(self, '打开文件夹', r'G:\Desktop\USMITK\data_image')

		# 设置路径
		self.lbl_open.setText(folder_dir)
		try:
			folder_name = folder_dir.split('/')[-1]
			folder_save = folder_dir.replace(folder_name, '')
		except:
			folder_name = 'demo'
			folder_save = '.'
		print(folder_save,folder_name)
		self.edt_save.setText('{}/{}.nii'.format(folder_save,folder_name))

		# 读取图片列表
		self.read_imgs(folder_dir)
		# self.img_list.setText('\n'.join(self.imgs_list))

		slm = QStringListModel()#创建mode
		slm.setStringList(self.imgs_list) #将数据设置到model
		self.img_list.setModel(slm)##绑定 listView 和 model
		return 


		# # 选中文件后，读出文件的内容，并设置成文本编辑框组件的显示文本
		# # folder_dir = QFileDialog.getOpenFileName(self, '打开文件', '/home')
		# if folder_dir[0]:
		# 	with open(folder_dir[0], 'r') as f:
		# 		data = f.read()
		# 		self.lbl_open.setText(data)


	imgs_list = []
	def read_imgs(self, path_folder=''):
		self.imgs_list = []
		for root, dirs, files in os.walk(path_folder):
			# 遍历文件
			for f in files:
				print('file:', os.path.join(root, f))
				if f.split('.')[-1].lower() in ['jpg', 'jpeg', 'bmp', 'png', 'gif', 'tif']:
				# if f.endswith('.jpg') or f.endswith('.png') or f.endswith('.bmp') or f.endswith('.jpeg'):
					self.imgs_list.append(os.path.join(root, f))
		return self.imgs_list

	def save_niis(self):
		# 判断保存合法性
		# self.imgs_list = glob(path_folder+'/*.jpg')
		assert(len(self.imgs_list)>0, "There are no images choosed")
		if self.imgs_list.__len__()<1:
			QMessageBox.about(self, '保存失败', '未选择图片，请选择图片文件夹')
			return 

		# 获取图片列表
		# print(self.imgs_list[0])
		imgs_2dus = [sitk.ReadImage(img) for img in self.imgs_list]
		imgs_2dus = [sitk.GetArrayFromImage(img) for img in imgs_2dus]
		imgs_2dus = [np.mean(img, axis=-1) for img in imgs_2dus]
		# print(self.imgs_list[0].shape)

		imgs_2dus = np.stack(imgs_2dus, axis=0)
		imgs_2dus = sitk.GetImageFromArray(imgs_2dus)

		# 保存文件
		path_save = self.edt_save.text()
		if not path_save.endswith('.nii'):
			path_save += '.nii'
		sitk.WriteImage(imgs_2dus, path_save)

		if os.path.exists(path_save):
			self.lbl_open.setText('成功保存到'+path_save)
			# # 信息框
			# QMessageBox.information(self, '框名', '内容', 按钮s, 默认按钮) 
			# # 问答框
			# QMessageBox.question(self, '框名', '内容', 按钮s, 默认按钮)
			# # 警告框
			# QMessageBox.warning(self, '框名', '内容', 按钮s, 默认按钮) 
			# # 危险框
			# QMessageBox.ctitical(self, '框名', '内容', 按钮s, 默认按钮)
			# 关于框
			# QMessageBox.about(self, '保存成功', '保存到'+path_save) 
		else:
			self.lbl_open.setText('保存失败!!!')
			# QMessageBox.about(self, '保存失败', '保存到'+path_save) 
		 


if __name__ == '__main__':
	app = QApplication(sys.argv)
	ex = FormatTookit()
	sys.exit(app.exec_())